package com.shujia.mllib

import org.apache.spark.ml.classification.LogisticRegressionModel
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.{DataFrame, SparkSession}

object Demo5UseModel {

  def main(args: Array[String]): Unit = {

    val spark: SparkSession = SparkSession
      .builder()
      .master("local[8]")
      .appName("image")
      .getOrCreate()
    import spark.implicits._

    //1、加载模型
    val model: LogisticRegressionModel = LogisticRegressionModel.load("spark/data/image_model")

    //单条预测
    //model.predict()
    //批量预测
    //model.transform()


    //读取需要识别的图片，转换成向量
    val testDF: DataFrame = spark.read
      .format("image")
      .load("E:\\data\\test")

    //编写自定义函数处理图片数据
    val comData: UserDefinedFunction = udf((data: Array[Byte]) => {
      val comData: Array[Double] = data.map(b => {
        //将图片中的像素点转换成0和1，白色:1,黑色：0
        if (b.toInt < 0) {
          1.0
        } else {
          0.0
        }
      })
      //转换成向量返回
      Vectors.dense(comData).toSparse
    })


    //将原始的图片转换成向量
    val testData: DataFrame = testDF.select($"image.origin", comData($"image.data") as "features")


    //使用模型识别图片
    val frame: DataFrame = model.transform(testData)

    frame
      .select($"origin", $"prediction")
      .show(false)

  }

}
